#!/usr/bin/env python

#based on flipper.py from Openeye

#convert smiles string to unique string
# keep or reject stereochemsitry information
#eg, unique smiles without chirality: True True False
#eg, "aromatic" smiles without chirality: True False False 


from openeye.oechem import *
import chemistry
import sys, string

if len(sys.argv) <> 5:
	print 'sterioisomerizer.py  <input_file> <output_file.smi> <flipall:True|False> <maxcenters>'
	sys.exit(1)

mol = OEGraphMol()

ifs = oemolistream(sys.argv[1])
out_file = open(sys.argv[2], 'w')
if sys.argv[3] == 'True':
	flip_all = True
else:
	flip_all = False
maxcenters = int(sys.argv[4])


#--------------------------------------------
def flip(mol, atom_list,asize, bond_list,bsize, stereoisomers):
  iso = True 
  kek = False
  chiral = True

  if asize>0:
    a=atom_list[asize-1]
    nbrs=[]
    for nbr in a.GetAtoms(): nbrs.append(nbr)

    a.SetStereo(nbrs, OEAtomStereo_Tetra, OEAtomStereo_Right)
    flip(mol,atom_list,asize-1,bond_list,bsize, stereoisomers)
    
    a.SetStereo(nbrs, OEAtomStereo_Tetra, OEAtomStereo_Left)
    flip(mol,atom_list,asize-1,bond_list,bsize, stereoisomers)
    return

  if bsize>0:
    b=bond_list[bsize-1]
    bgn = b.GetBgn()
    end = b.GetEnd()

    nbrs=[]
    for nbr in bgn.GetAtoms():
      if nbr.GetIdx() != end.GetIdx(): 
          nbrs.append(nbr)
          break
       
    for nbr in end.GetAtoms():
      if nbr.GetIdx() != bgn.GetIdx(): 
          nbrs.append(nbr)
          break
        
    b.SetStereo(nbrs, OEBondStereo_CisTrans, OEBondStereo_Cis)
    flip(mol,atom_list,asize,bond_list,bsize-1, stereoisomers)
    b.SetStereo(nbrs, OEBondStereo_CisTrans, OEBondStereo_Trans)
    flip(mol,atom_list,asize,bond_list,bsize-1, stereoisomers)
    return
    
  smi = chemistry.CanSmi(mol,iso,kek, chiral)
  if smi not in stereoisomers:
  	stereoisomers.append(smi)
  return stereoisomers
#--------------------------------------------------
if __name__=="__main__":

  iso = True 
  kek = False
  chiral = True

  # build atom predicate
  isChiralAtom = OEIsChiralAtom()
  if flip_all:
    okAtom = isChiralAtom
  else:
    atomStereoSpecified = OEHasAtomStereoSpecified()
    notAtomStereoSpecified = OENotAtom(atomStereoSpecified)
    okAtom = OEAndAtom(isChiralAtom,notAtomStereoSpecified)
  
  # build bond predicate
  isChiralBond = OEIsChiralBond()
  if flip_all:
    okBond = isChiralBond
  else:
    bondStereoSpecified = OEHasBondStereoSpecified()
    notBondStereoSpecified = OENotBond(bondStereoSpecified)
    okBond = OEAndBond(isChiralBond,notBondStereoSpecified)
    #what about rings? => they are not counted (that's is correct)
  
  for mol in ifs.GetOEGraphMols():

    OEPerceiveChiral(mol)
    stereoisomers = []    
    nchia=0
    nchib=0
    atom_list=[]
    bond_list=[]
    for atom in mol.GetAtoms(okAtom):
      nchia+=1
      atom_list.append(atom)

    for bond in mol.GetBonds(okBond):
      nchib+=1
      bond_list.append(bond)

    #print nchia, nchib

    if nchia+nchib == 0:
	stereoisomers.append(chemistry.CanSmi(mol,iso,kek, chiral))
	s_type = 'specified'
    elif nchia+nchib > maxcenters:
	stereoisomers.append(chemistry.CanSmi(mol,iso,kek, chiral))
	s_type = 'exceeded max centres (' + str(maxcenters) + ')'
      
    else:
	
	flip(mol,atom_list,len(atom_list),bond_list,len(bond_list), stereoisomers)
	s_type = 'guessed'


    for moli in stereoisomers:
	out_file.write(moli + '\t' + mol.GetTitle() + '\t' + s_type + '\n')

out_file.close()

print 'done'




